import socket
import os
import sys
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad


# 密钥和IV
FIXED_KEY = b"MySecretKey12345"
FIXED_IV = b"InitVector123456"


# 加密函数
def encrypt_data(data):
    cipher = AES.new(FIXED_KEY, AES.MODE_CBC, FIXED_IV)
    padded_data = pad(data, AES.block_size)
    encrypted_data = cipher.encrypt(padded_data)
    return encrypted_data


# 发送文件函数
def send_file(file_path):
    # 检查文件是否存在
    if not os.path.exists(file_path):
        print(f"错误: 文件不存在: {file_path}")
        return

    # 创建套接字
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    # 连接服务器
    host = 'localhost'
    port = 9999
    client_socket.connect((host, port))
    print(f"已连接到服务器: {host}:{port}")

    # 发送文件名
    filename = os.path.basename(file_path)
    client_socket.send(filename.encode('utf-8'))

    # 读取文件
    with open(file_path, 'rb') as f:
        file_data = f.read()

    # 加密文件
    print("正在加密...")
    encrypted_data = encrypt_data(file_data)

    # 发送文件大小
    client_socket.send(str(len(encrypted_data)).encode('utf-8'))

    # 发送加密数据
    client_socket.sendall(encrypted_data)
    print("文件发送完成")

    # 接收服务器响应
    response = client_socket.recv(1024).decode('utf-8')
    print(f"服务器回复: {response}")

    # 关闭连接
    client_socket.close()
    print("连接已关闭")


# 程序入口
if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("用法: python simple_aes_client.py <文件路径>")
    else:
        file_path = sys.argv[1]
        send_file(file_path)